#include "treebp.h"
#include <time.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif 

#define ERRIDX    -1
#define MINDIFF   0.00001
#define CHKFREQ   100

clock_t mystart, myend;
double elapsed;

typedef double (* deflocalf)(node_t *, int, double *);
typedef double (* defpairf)(node_t *, node_t *, int, int, double *);

double getlocal(node_t *pnode, int yi, double *para)
{
   nodehash_t *tmphash;
   node_t *pnbr;
   double val = 0.0;

   val = pnode->val[0];
   for(tmphash=pnode->nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
      pnbr = tmphash->key;
      /*val += fabs(pnode->val[0]-pnbr->val[0]);*/
      val += pnbr->val[0];
   }
   val = fabs(val/(pnode->nbrs+1) - yi);
   
   return val;
}

double getpair(node_t *pnode, node_t *pnbr, int ys, int yt, double *para)
{
   if(ys == yt)
      return 0.0;
   else
      return 1.0;
}

int chkconverge(node_t *phead, int *pys, int totalnodes)
{
   int i;
   double diffs = 0;

   for(i=0;i<totalnodes;i++) {
      diffs += abs(phead[i].ymax - pys[i]);
   }
   if(diffs < MINDIFF)
      return 1;
   else
      return 0;
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
   int         ifchk;
   int         i, j, outdim[2], boxid = 0;
   int         totalnodes=0, maxnbrs=0, maxxdims=0, maxloop = 0;
   int         *pnbrnum, **pnbridx, *p2oldys, pys[NUMOFSTATE];
   double      *p4inputnbrnum, *p4inputidx, *p4inputxval, *p4inputys;
   double      **pxval, para1[1], para2[1], pk;
   node_t      *phead;
   char        *pout;

   /* inputs: number of total nodes, max number of neighbors for a node, 
    * max dimension of the data associated with a node, the value of state
    * variable y, array of neighbor number for each node, neighbor index for 
    * each node, observation data vector for each node, max number of loop for
    * message passing
    */
   totalnodes     = (int) mxGetScalar(prhs[0]);
   maxnbrs        = (int) mxGetScalar(prhs[1]);
   maxxdims       = (int) mxGetScalar(prhs[2]);
   p4inputys      = mxGetPr(prhs[3]);
   p4inputnbrnum  = mxGetPr(prhs[4]);
   p4inputidx     = mxGetPr(prhs[5]);
   p4inputxval    = mxGetPr(prhs[6]);
   maxloop        = (int) mxGetScalar(prhs[7]);
   
   if(nrhs != 8) 
      mexErrMsgTxt("Wrong number of input parameters!\n");

   para1[0] = 1.0;
   para2[0] = 1.0;
   p2oldys = (int *) mxMalloc(sizeof(int)*totalnodes);
   memset((char *)p2oldys, 0, sizeof(int)*totalnodes);
   /* allocate memory for neighbor index and observed value x
    */
   pnbrnum = (int *) mxMalloc(sizeof(int)*totalnodes);
   pnbridx = (int **) mxMalloc(sizeof(int *)*totalnodes);
   for(i=0;i<totalnodes;i++) {
      pnbridx[i] = (int *) mxMalloc(sizeof(int)*maxnbrs);
   }
   pxval = (double **) mxMalloc(sizeof(double *)*totalnodes);
   for(i=0;i<totalnodes;i++) {
      pxval[i] = (double *) mxMalloc(sizeof(double)*maxxdims);
   }

   /* reassign neighbor index and observed value
    * could be improved for better mem usage but convenient for now
    */
   mystart = clock();
   for(i=0;i<NUMOFSTATE;i++) {
      pys[i] = (int) p4inputys[i];
   }

   for(i=0;i<totalnodes;i++) {
      pnbrnum[i] = (int) p4inputnbrnum[i];
      for(j=0;j<maxnbrs;j++) {
         pnbridx[i][j] = (int) p4inputidx[i*maxnbrs+j];
      }
      for(j=0;j<maxxdims;j++) {
         pxval[i][j] = p4inputxval[i*maxxdims+j];
      }
   }
   myend = clock();
   elapsed = ((double)(myend-mystart))/CLOCKS_PER_SEC;
   printf("Beginning assignment time is %f\n", elapsed);

   /* initilization of all nodes
    */
   mystart = clock();
   phead = initall(totalnodes, pxval, pys, pnbrnum, pnbridx);
   myend = clock();
   elapsed = ((double)(myend-mystart))/CLOCKS_PER_SEC;
   printf("nodes init time is %f\n", elapsed);

   /* assign function pointer of each node
    */
   mystart = clock();
   for(i=0;i<totalnodes;i++) {
      phead[i].clocal = getlocal;
      phead[i].cpair = getpair;
   }

   boxid = 1;
   ifchk = 1;
   for(i=0;i<maxloop;i++) {
      boxid = 1 - boxid;
      if((i % 2) == 0){
         printf("loop#%d: run forward max-msg passing\n", i);
         fflush(NULL);
         runmaxbpfd(phead, totalnodes, para1, para2, boxid);
      } else {
         printf("loop#%d, run backward max-msg passing\n", i);
         fflush(NULL);
         runmaxbpbd(phead, totalnodes, para1, para2, boxid);
      }
      if((ifchk % CHKFREQ) == 0) {
         bdbt4bp(phead, totalnodes, boxid);
         if(chkconverge(phead, p2oldys, totalnodes)) {
            break;
         } else {
            for(j=0;j<totalnodes;j++) {
               p2oldys[j] = phead[j].ymax;
            }
         }
      }
      ifchk += 1;
   }

   boxid = 1 - boxid;
   for(i=0;i<totalnodes;i++) {
      computebi(&phead[i], para1, &pk, boxid);
   }
   myend = clock();
   elapsed = ((double)(myend-mystart))/CLOCKS_PER_SEC;
   printf("loop time is %f\n", elapsed);

   mystart = clock();

#ifdef USEBACKTRACE
   bdbt4bp(phead, totalnodes, boxid);
#elif USEFORWORDTRACE 
   fdbt4bp(phead, totalnodes, boxid);
#else
   sortbybi(phead, totalnodes);
#endif
   myend = clock();
   elapsed = ((double)(myend-mystart))/CLOCKS_PER_SEC;
   printf("backtrace time is %f\n", elapsed);


   outdim[0] = totalnodes;
   outdim[1] = 1;
   plhs[0] = mxCreateNumericArray(1,outdim,mxINT8_CLASS,mxREAL);
   pout = (char *) mxGetPr(plhs[0]);
   for(i=0;i<totalnodes;i++) {
      pout[i] = (unsigned char) phead[i].ymax;
   }
 
   mystart = clock();
   for(i=0;i<totalnodes;i++) {
      mxFree(pnbridx[i]);
      mxFree(pxval[i]);
   }
   mxFree(pnbridx);
   mxFree(pxval);
   mxFree(pnbrnum);
   mxFree(p2oldys);
   freeallnodes(phead, totalnodes);
   myend = clock();
   elapsed = ((double)(myend-mystart))/CLOCKS_PER_SEC;
   printf("release mem time is %f\n", elapsed);

}
